import math

from strange_io import SIGNATURE, _load_bytearray, get_format

class StrangeVoxels:
    def __init__(self, width=20, height=20, depth=20, palette_size:int=1):
        self.palette_size:int = palette_size
        self.height:int = height
        self.width:int = width
        self.depth:int = depth
        self.voxels:list[list[list[int]]] = [[[0 for _ in range(depth)] for _ in range(height)] for _ in range(width)]

    def dump(self) -> dict:
        out = {}
        out['palette_size'] = self.palette_size
        out['height'] = self.height
        out['width'] = self.width
        out['depth'] = self.depth
        out['model'] = self.voxels
        return out

    def load(self, model) -> dict:
        self.palette_size = model['palette_size']
        self.height = model['height']
        self.width = model['width']
        self.depth = model['depth']
        self.voxels = model['model']

    def dump_bytes(self) -> bytes:
        return self.encode(self.dump())

    def load_bytes(self, model:bytes):
        self.load(self.decode(model))

    def dump_file(self, path):
        """ Dump document to file """
        out = bytearray()
        header = bytearray(SIGNATURE + b'voxels\n')
        out += header
        out += self.dump_bytes()
        with open(path, 'wb') as f:
            f.write(out)

    def load_file(self, path, ignore_type=False) -> dict:
        """ Load document from file """
        if get_format(path) == b'voxels' or ignore_type:
            file = _load_bytearray(path, remove_type=True)
            file += bytearray([0x00]) # Add a byte at the end of the file
            self.load_bytes(file)
        else:
            raise TypeError('Attempted to load non-voxels document as voxels document')

    @staticmethod
    def encode(img:dict) -> bytes:
        out = bytearray()
        out += dynamic_integer.encode(img['width'])[0]
        out += dynamic_integer.encode(img['height'])[0]
        out += dynamic_integer.encode(img['depth'])[0]
        for plane in img['model']:
            for column in plane:
                for voxel in column:
                    out += voxel.to_bytes(1, 'big')
        out += img['palette_size'].to_bytes(1, 'big')
        return out

    @staticmethod
    def decode(file:bytes) -> dict:
        out = {}
        out['palette_size'] = file[len(file)-1]
        #region width
        width = dynamic_integer.decode(file)[0]
        out['width'] = width
        #endregion
        #region height
        height = dynamic_integer.decode(file)[0]
        out['height'] = height
        #endregion
        #region depth
        depth = dynamic_integer.decode(file)[0]
        out['depth'] = depth
        #endregion
        out['model'] = []
        column = []
        plane = []
        for voxel in file:
            if len(column) == out['height']:
                plane.append([*column])
                column.clear()
            if len(plane) == out['depth']:
                out['model'].append([*plane])
                column.clear()
            column.append(voxel)
        return out
